# visualization.py
import matplotlib.pyplot as plt
from typing import Dict, Any

import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from FastAPI.models import PlannerConfig
from algorithm.Astar import AStarPlanner

def plot_trajectories(response_data: Dict[str, Any], config_data: Dict[str, Any], save_image: bool = True):
    """
    根据API返回的轨迹数据和配置数据进行绘图。

    :param response_data: API响应中的 'data' 字段，包含 trajectories 和 time_paths
    :param config_data: 发送给API的 'config' 字段
    """
    time_paths = response_data.get('time_paths', {})
    if not time_paths:
        print("没有可供可视化的轨迹数据。")
        return

    # 使用配置重建一个临时的planner以获取坐标转换方法和地图信息
    config = PlannerConfig.model_validate(config_data)
    min_x, min_y, max_x, max_y = config.map_bounds
    obs_x = [min_x, max_x, min_x, max_x]
    obs_y = [min_y, min_y, max_y, max_y]
    
    # 注意：这里的 fine_grid_size 应该与规划时使用的相匹配
    temp_planner = AStarPlanner(obs_x, obs_y, config.fine_grid_size, config.robot_radius)
    
    all_paths_list = list(time_paths.values())
    uav_ids = list(time_paths.keys())
    
    n = len(all_paths_list)
    if n == 0:
        print("没有路径可以绘制。")
        return
        
    colors = plt.get_cmap('tab10')
    fig, ax = plt.subplots(figsize=(12, 12))

    # 绘制轨迹
    for idx, path in enumerate(all_paths_list):
        if not path: continue
        
        uav_id = uav_ids[idx]
        color = colors(idx % 10)
        
        xs = [temp_planner.calc_grid_position(xi, temp_planner.min_x) for xi, _, _ in path]
        ys = [temp_planner.calc_grid_position(yi, temp_planner.min_y) for _, yi, _ in path]
        
        ax.plot(xs, ys, '-', linewidth=2, color=color, label=f'UAV {uav_id} Trajectory')
        ax.scatter(xs[0], ys[0], marker='o', s=100, color=color, edgecolors='k', zorder=5, label=f'UAV {uav_id} Start')
        ax.scatter(xs[-1], ys[-1], marker='X', s=100, color=color, edgecolors='k', zorder=5, label=f'UAV {uav_id} Goal')

    ax.set_title("UAV Planned Trajectories")
    ax.set_xlabel("X coordinate (m)")
    ax.set_ylabel("Y coordinate (m)")
    ax.grid(True, linestyle='--', alpha=0.6)
    ax.set_aspect('equal', adjustable='box')
    
    # 整理图例
    handles, labels = ax.get_legend_handles_labels()
    by_label = dict(zip(labels, handles))
    ax.legend(by_label.values(), by_label.keys())
    
    # 根据参数决定是否保存图片
    if save_image:
        vis_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'visualization')
        os.makedirs(vis_dir, exist_ok=True)
        trajectory_image_path = os.path.join(vis_dir, 'trajectory_visualization.png')
        plt.savefig(trajectory_image_path, dpi=300, bbox_inches='tight')
        print(f"轨迹可视化图片已保存到: {trajectory_image_path}")
    
    plt.show()

